#include "kernel_operator.h"
#include <cstdlib>
using namespace AscendC;
constexpr int32_t BUFFER_NUM = 2;

class KernelTriu
{
public:
    __aicore__ inline KernelTriu() {}
    __aicore__ inline void Init(GM_ADDR x, GM_ADDR y,
                                uint32_t ALIGN_NUM, uint32_t unitCount, uint32_t totalLength,
                                uint32_t tileNum, uint32_t tileLength, uint32_t lasttileLength,
                                uint32_t h, uint32_t w, int32_t diagonal)
    {
        this->ALIGN_NUM = ALIGN_NUM;
        this->totalLength = totalLength;
        this->tileNum = tileNum;
        this->tileLength = tileLength;
        this->lasttileLength = lasttileLength;
        this->w = w;
        this->wAligned = (w + ALIGN_NUM - 1) / ALIGN_NUM * ALIGN_NUM;
        this->h = h;
        this->unitCount = unitCount;
        this->diagonal = diagonal;
        

        xGm.SetGlobalBuffer((__gm__ DTYPE_X *)x, totalLength * unitCount);
        yGm.SetGlobalBuffer((__gm__ DTYPE_Y *)y, totalLength * unitCount);


        pipe.InitBuffer(xQueue, BUFFER_NUM, tileLength * sizeof(DTYPE_X)); //自动向上32Byte对齐
        pipe.InitBuffer(yQueue, BUFFER_NUM, tileLength * sizeof(DTYPE_Y));
        
        pipe.InitBuffer(Tmp_F, tileLength * sizeof(float));
        pipe.InitBuffer(Tmp, tileLength * sizeof(int32_t));

        pipe.InitBuffer(Row, tileLength * sizeof(int32_t));
        pipe.InitBuffer(Col, tileLength * sizeof(int32_t));
        pipe.InitBuffer(Gidx, tileLength * sizeof(int32_t));

        pipe.InitBuffer(Bit, tileLength * sizeof(uint8_t));

        pipe.InitBuffer(Row_F, tileLength * sizeof(float));
        pipe.InitBuffer(Col_F, tileLength * sizeof(float));
        pipe.InitBuffer(Lidx_F, tileLength * sizeof(float));
        pipe.InitBuffer(Gidx_F, tileLength * sizeof(float));
    }


    __aicore__ inline void InitIdx()
    {
        LocalTensor<float> lidx_f = Lidx_F.Get<float>();
        Muls(lidx_f, lidx_f, (float)0.0, tileLength);   //0 0 0...
        for(int32_t i = 1; i < tileLength; i++)
        {
            lidx_f.SetValue(i, (float)i);    //0 1 2 ....
        }
    }
    
    __aicore__ inline void Process()
    {
        InitIdx();  //初始化idx
        uint32_t  loopCount = tileNum;
        for(int j=0; j<unitCount; j++)
        {
            for (int32_t i = 0; i < loopCount - 1; i++) 
            {                
                CopyIn(j, i, tileLength);
                Compute(j, i, tileLength);
                CopyOut(j, i, tileLength);
            }

            CopyIn(j, tileNum - 1, lasttileLength);
            Compute(j, tileNum - 1, lasttileLength);
            CopyOut(j, tileNum - 1, lasttileLength);
        }        
    }

private:
    __aicore__ inline void CopyIn(int32_t j, int32_t progress, uint32_t length)
    {
        LocalTensor<DTYPE_X> xLocal = xQueue.AllocTensor<DTYPE_X>();
        uint32_t startindex = j * totalLength + progress * tileLength;
        DataCopy(xLocal, xGm[startindex], length);
        xQueue.EnQue(xLocal);
    }

    __aicore__ inline void Compute(int32_t j, int32_t progress, uint32_t length)
    {
        LocalTensor<DTYPE_X> xLocal = xQueue.DeQue<DTYPE_X>();
        LocalTensor<DTYPE_Y> yLocal = yQueue.AllocTensor<DTYPE_Y>();

        LocalTensor<uint8_t> bit = Bit.Get<uint8_t>();

        LocalTensor<float> tmp_f = Tmp_F.Get<float>();
        LocalTensor<int32_t> tmp = Tmp.Get<int32_t>();
        
        LocalTensor<int32_t> row = Row.Get<int32_t>();
        LocalTensor<int32_t> col = Col.Get<int32_t>();

        LocalTensor<float> row_f = Row_F.Get<float>();
        LocalTensor<float> col_f = Col_F.Get<float>();

        LocalTensor<float> lidx_f = Lidx_F.Get<float>();
        LocalTensor<float> gidx_f = Gidx_F.Get<float>();
        LocalTensor<int32_t> gidx = Gidx.Get<int32_t>();
        
        Adds(gidx_f, lidx_f, (float)progress * tileLength, length);    //每个batch内部当前tiling的下标
        Cast(gidx, gidx_f, RoundMode::CAST_FLOOR, length);   //int32 的gidx

        Muls(tmp_f, gidx_f, (float)1.0/(float)w, length);    //row 的float
        Cast(row, tmp_f, RoundMode::CAST_FLOOR, length); //row 为int32 类型，表示每个元素所在的 row
        
        Muls(tmp, row, w, length);
        Sub(col, gidx, tmp, length); //col 的 int32

        Cast(row_f, row, RoundMode::CAST_NONE, length);   //int32->float
        Cast(col_f, col, RoundMode::CAST_NONE, length);

        Adds(row_f, row_f, (float)diagonal, length);
        Compare(bit, col_f, row_f, CMPMODE::GE, length);
        Select(yLocal, bit, xLocal, (DTYPE_X)0.0, SELMODE::VSEL_TENSOR_SCALAR_MODE, length);

        xQueue.FreeTensor(xLocal);
        yQueue.EnQue<DTYPE_Y>(yLocal);
    }

    __aicore__ inline void CopyOut(int32_t j, int32_t progress, uint32_t length)
    {
        LocalTensor<DTYPE_Y> yLocal = yQueue.DeQue<DTYPE_Y>();
        uint32_t startindex = j * totalLength + progress * tileLength;
        DataCopy(yGm[startindex], yLocal, length);
        yQueue.FreeTensor(yLocal);
    }


    TPipe pipe;
    TQue<QuePosition::VECIN, BUFFER_NUM> xQueue;
    TQue<QuePosition::VECOUT, BUFFER_NUM> yQueue;

    GlobalTensor<DTYPE_X> xGm;
    GlobalTensor<DTYPE_Y> yGm;
    TBuf<QuePosition::VECCALC> Tmp, Row, Col, Gidx, Bit;
    TBuf<QuePosition::VECCALC> Tmp_F, Row_F, Col_F, Lidx_F, Gidx_F;

    uint32_t ALIGN_NUM;
    uint32_t totalLength;
    uint32_t tileLength;
    uint32_t tileNum;
    uint32_t lasttileLength;
    int32_t h;
    int32_t w;
    int32_t wAligned;
    int32_t unitCount;
    int32_t diagonal;
};


extern "C" __global__ __aicore__ void triu(GM_ADDR x, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling) {
    GET_TILING_DATA(tiling_data, tiling);
    // TODO: user kernel impl

    KernelTriu op;
    op.Init(x, y, tiling_data.ALIGN_NUM, tiling_data.unitCount, tiling_data.totalLength,
                            tiling_data.tileNum, tiling_data.tileLength, tiling_data.lasttileLength,
                            tiling_data.h, tiling_data.w, tiling_data.diagonal);
    op.Process();
}